# -*- coding: utf-8 -*-
"""\
This file is part of the ptycho package.

Utility functions and classes to support MPI computing.
"""
__all__ = ['MPIenabled', 'size', 'rank', 'comm', 'MPI']

from utils import verbose
global size, rank

try:
    from mpi4py import MPI
    comm = MPI.COMM_WORLD
    size = comm.Get_size()
    rank = comm.Get_rank()
except:
    verbose(2, 'MPI initialization failed. Proceeding with one processor')
    MPI = None
    size = 1
    rank = 0
    comm = None

MPIenabled = not (size==1)

def useMPI(do=None):
    """\
    Toggle using MPI or not.    
    """
    global MPIenabled
    if do is None:
        return MPIenabled
    if MPI is None:
        MPIenabled = False
    else:
        MPIenabled = do
